#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <algorithm>

enum { kRed, kBlack };

class Node {
public:
  Node(int key, Node *parent) : 
      left(NULL), right(NULL), parent(parent), key(key), color(kRed) {
    left = this->NullNode();
    right = this->NullNode();
  }

  Node *Uncle() { return parent->Sibling(); }
  Node *GrandPa() { return parent->parent; }
  Node *Sibling() {
    if (this == parent->left)
      return parent->right;
    else
      return parent->left;
  }
  Node *Successor();
  Node *Predecessor();
  void Copy(const Node &n);

  Node *left, *right, *parent;
  int key;
  char color;

  bool IsNull() const { return this == this->NullNode(); }
  static Node *NullNode() { return &null_node; }

private:
  Node() : left(NULL), right(NULL), parent(NULL), key(0), color(kBlack) { }

  static Node null_node;
};

Node Node::null_node;

Node *Node::Successor() {
  Node *p = this->right;
  while (!p->left->IsNull()) p = p->left;
  return p;
}

Node *Node::Predecessor() {
  Node *p = this->left;
  while (!p->right->IsNull()) p = p->left;
  return p;
}

void Node::Copy(const Node &n) {
  this->key = n.key;
}

Node *Search(Node *root, int key) {
  Node *p = root;
  while (!p->IsNull()) {
    if (p->key < key)
      p = p->right;
    else if (p->key > key)
      p = p->left;
    else return p;
  }
  return p;
}

void InsertCase1(Node **root, Node *n);
void InsertCase2(Node **root, Node *n);
void InsertCase3(Node **root, Node *n);
void InsertCase4(Node **root, Node *n);
void InsertCase5(Node **root, Node *n);

void DeleteCase1(Node **root, Node *n);
void DeleteCase2(Node **root, Node *n);
void DeleteCase3(Node **root, Node *n);
void DeleteCase4(Node **root, Node *n);
void DeleteCase5(Node **root, Node *n);
void DeleteCase6(Node **root, Node *n);

void Replace(Node **root, Node *n, Node *p) {
  if (n == *root) {
    *root = p;
  }
  else {
    if (n == n->parent->left) n->parent->left = p;
    else n->parent->right = p;
  }
  p->parent = n->parent;
}

void RotateLeft(Node **root, Node *n) {  
  //printf("RotateLeft\n");
  Node *r = n->right;  

  Replace(root, n, r);  
  
  n->right = r->left;
  if (!r->left->IsNull()) r->left->parent = n;

  r->left = n;
  n->parent = r;  
}

void RotateRight(Node **root, Node *n) {  
  //printf("RotateRight\n");
  Node *l = n->left;  

  Replace(root, n, l);  

  n->left = l->right;
  if (!l->right->IsNull()) l->right->parent = n;

  l->right = n;
  n->parent = l;  
}

// n is red
void InsertCase1(Node **root, Node *n) {
  if (n->parent->IsNull()) {
    //printf("case1\n");
    n->color = kBlack;
    *root = n;
  }
  else {
    InsertCase2(root, n);
  }
}

// n->parent is not NULL, thus is not root
void InsertCase2(Node **root, Node *n) {
  if (n->parent->color == kBlack) {
    //printf("case2\n");
    return;
  }
  else {
    InsertCase3(root, n);
  }
}

// n->parent is red
void InsertCase3(Node **root, Node *n) {
  if (n->Uncle()->color == kRed) {
    //printf("case3\n");
    n->parent->color = kBlack;
    n->Uncle()->color = kBlack;
    n->GrandPa()->color = kRed;
    InsertCase1(root, n->GrandPa());
  }
  else {
    InsertCase4(root, n);
  }
}

// !!!case4 and case5 can be merged to promote efficiency

// n->parent is red, n->Uncle() is black, or NULL
void InsertCase4(Node **root, Node *n) { 
  if (n == n->parent->right && n->parent == n->GrandPa()->left) {
    //printf("case4\n");
    RotateLeft(root, n->parent);
    n = n->left;
  }
  else if (n == n->parent->left && n->parent == n->GrandPa()->right) {
    //printf("case4\n");
    RotateRight(root, n->parent);
    n = n->right;
  }
  InsertCase5(root, n);
}

// case 1: n->parent is a left son, n is a left son;
// case 2: n->parent is a right son, n is a right son;
void InsertCase5(Node **root, Node *n) {
  //printf("case5\n");
  n->parent->color = kBlack;
  n->GrandPa()->color = kRed;
  if (n == n->parent->left)
    RotateRight(root, n->GrandPa());
  else
    RotateLeft(root, n->GrandPa());
}

void Insert(Node **root, int key) {
  Node **n = root;
  Node *parent = *n;
  while (!(*n)->IsNull()) {
    parent = *n;
    if (parent->key < key)
      n = &parent->right;
    else
      n = &parent->left;
  }
  *n = new Node(key, parent);
  //printf("Inserted\n");
  InsertCase1(root, *n);
}

void DeleteCase6(Node **root, Node *n) {  
  Node *s = n->Sibling();
 
  s->color = n->parent->color;
  n->parent->color = kBlack;

  if (n == n->parent->left) {
    s->right->color = kBlack;
    RotateLeft(root, n->parent);
  }
  else {
    s->left->color = kBlack;
    RotateRight(root, n->parent);
  }
}

void DeleteCase5(Node **root, Node *n) { 
  if (n == n->parent->left) {
    Node *s = n->parent->right;
    if (s->right->color == kBlack &&
        s->left->color == kRed) { 
      s->color = kRed;
      s->left->color = kBlack;      
      RotateRight(root, s);
    }
  }
  else {
    Node *s = n->parent->left;
    if (s->left->color == kBlack &&
        s->right->color == kRed) { 
      s->color = kRed;
      s->right->color = kBlack;
      RotateLeft(root, s);
    }
  }  
  DeleteCase6(root, n);
}

void DeleteCase4(Node **root, Node *n) { 
  Node *s = n->Sibling();
  if (n->parent->color == kRed &&
      s->color == kBlack &&
      s->left->color == kBlack &&
      s->right->color == kBlack) {
    s->color = kRed;
    n->parent->color = kBlack;
  }
  else {
    DeleteCase5(root, n);
  }
}

void DeleteCase3(Node **root, Node *n) { 
  Node *s = n->Sibling();
  if (n->parent->color == kBlack &&
      s->color == kBlack &&
      s->left->color == kBlack &&
      s->right->color == kBlack) {
    s->color = kRed;
    DeleteCase1(root, n->parent);
  }
  else {
    DeleteCase4(root, n);
  }
}

// transform sibling color to black
void DeleteCase2(Node **root, Node *n) {   
  Node *s = n->Sibling();
  if (s->color == kRed) {
    n->parent->color = kRed;
    s->color = kBlack;
    if (n == n->parent->left)
      RotateLeft(root, n->parent);
    else
      RotateRight(root, n->parent);
  }
  DeleteCase3(root, n);
}

void DeleteCase1(Node **root, Node *n) {  
  if (n->parent->IsNull()) return;
  DeleteCase2(root, n);
}

void Delete(Node **root, Node *n) {
  Node *child = n->left->IsNull() ? n->right : n->left;
  Replace(root, n, child);
  if (n->color == kBlack) {
    if (child->color == kRed) {
      child->color = kBlack;
    }
    else {
      DeleteCase1(root, child);
    }
  }  
  delete n;
}

void Delete(Node **root, int key) {
  Node *n = Search(*root, key);
  if (n->IsNull()) return;

  if (!n->left->IsNull() && !n->right->IsNull()) {
    Node *p = n->Successor();
    n->Copy(*p);
    n = p;
  }
  Delete(root, n);
}

void DeleteTree(Node **root) {
  if ((*root)->IsNull()) return;
  DeleteTree(&(*root)->left);
  DeleteTree(&(*root)->right);
  delete *root;
  *root = Node::NullNode();
}

class Tree {
public:
  Tree() : root(Node::NullNode()) { }
  ~Tree() { DeleteTree(&root); }

  void Insert(int key) { ::Insert(&root, key); }
  void Insert(int keys[], int n) {
    for (int i = 0; i < n; ++i) this->Insert(keys[i]);
  }

  void Delete(int key) { ::Delete(&root, key); }
  void Delete(int keys[], int n) {
    for (int i = 0; i < n; ++i) this->Delete(keys[i]);
  }

  Node *Search(int key) { return ::Search(root, key); }

  bool IsValid() {
    if (root->IsNull()) {
      printf("Empty\n");
      return false;
    }
    if (root->color != kBlack) {
      printf("root is not black\n");
      return false;
    }
    int n_black = -1;
    int n_current_black = 0;
    return Check(root, n_current_black, &n_black);
  }

  bool Check(Node *n, int n_current_black, int *n_black) {
    if (n->IsNull()) {
      ++n_current_black;
      if (*n_black < 0) *n_black = n_current_black;
      else if (*n_black != n_current_black) {
        printf("different number of black nodes in path: %d, %d\n",
               *n_black, n_current_black);
        return false;
      }
      return true;
    }
    if (n->color == kRed) {
      if (!((n->left->IsNull() || n->left->color == kBlack) &&
            (n->right->IsNull() || n->right->color == kBlack))) {
        printf("not all sons of red node are black\n");
        return false;
      }
    }
    else {
      ++n_current_black;
    }

    if ((!n->left->IsNull() && n->left->key > n->key) ||
        (!n->right->IsNull() && n->right->key < n->key)) {
      printf("not a valid binary search tree\n");
      return false;
    }
    if (!Check(n->left, n_current_black, n_black)) return false;
    if (!Check(n->right, n_current_black, n_black)) return false;

    return true;
  }
  
  Node *root;
};

int *RandInt(int n) {
  printf("n = %d\n", n);

  int *data = new int[n];
  for (int i = 0; i < n; ++i) {
    data[i] = rand() % n;
  }
  printf("data ok!\n");
  return data;
}

int Search(const int A[], int n, int key) {
  for (int i = 0; i < n; ++i) {
    if (A[i] == key) return i;
  }
  return -1;
}

int *GenTestSamples(const int A[], int n, int m) {
  int *T = new int[m];
  for (int i = 0; i < m; ++i) {
    if (rand() % 1 == 0) {
      T[i] = rand() % n;
    }
    else {
      T[i] = A[rand() % n];
    }
  }
  printf("%d search samples\n", m);
  return T;
}

void SearchTest(const int A[], int n, Tree &tree, const int T[], int m) {
  for (int i = 0; i < m; ++i) {
    Node *p = tree.Search(T[i]);
    int j = Search(A, n, T[i]);
    if (!((p->IsNull() && j < 0) ||
        (!p->IsNull() && j >= 0))) {
      printf("Search Error, there may be problem in the tree\n");
      break;
    }
  }
}

void SearchTest(const int A[], int n, Tree &tree) {
  int m = std::max(10, n / 100);
  int *T = GenTestSamples(A, n, m);

  SearchTest(A, n, tree, T, m);

  delete[] T;
}

void SearchTest(Tree &tree, const int T[], int m) {
  int count = 0;
  for (int i = 0; i < m; ++i)
    if (!tree.Search(T[i])->IsNull()) ++count;
  printf("find %d samples in tree\n", count);
}

void DeleteTest(const int A[], int n, Tree &tree) {
  int m = std::max(10, n / 100);
  int *T = GenTestSamples(A, n, m);

  SearchTest(A, n, tree, T, m);  
  
  SearchTest(tree, T, m);
  tree.Delete(T, m);
  tree.IsValid();

  SearchTest(tree, T, m);
  tree.Delete(T, m);
  tree.IsValid();  

  SearchTest(tree, T, m);
  tree.Delete(T, m);
  tree.IsValid();

  SearchTest(tree, T, m);

  delete[] T;
}

int main(int argc, const char *argv[]) {
  if (argc < 2) {
    printf("Usage: <cmd> <number of nodes>\n");
    return -1;
  }

  srand(time(NULL));
  int n = std::max(2, rand() % atoi(argv[1]));

  int *A = RandInt(n);
  Tree tree;
 
  tree.Insert(A, n);
  if (tree.IsValid()) {
    printf("Tree is valid\n");
  }  

  //SearchTest(A, n, tree);
  DeleteTest(A, n, tree);

  delete[] A;
  return 0;
}

